import numpy as np
import cv2
import math

from numpy import ma
from EdgeDetect import * 
def getKeyPoints(path):

    img = cv2.imread(path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    sift = cv2.SIFT_create()
    kps = sift.detect(gray, None)
    # print(len(kps))
    img=cv2.drawKeypoints(gray,kps,img)#绘制关键点
    # print(kps[0].pt)
    # cv2.imshow('sp',img)
    cv2.imwrite("./SIFT/keyPoints_rawImage_rotate.jpg", img)
    # cv2.waitKey(0)

    keypoints = []
    for kp in kps:
        x, y = int(kp.pt[0]), int(kp.pt[1])
        if x >= 8 and x < 248 and y >= 8 and y < 248:
            keypoints.append((x, y))
    print(len(keypoints))
    return gray, keypoints

def getGrad(image):
    h = image.shape[0]
    w = image.shape[1]
    print(image.shape)
    # xtmp = [[-1,0,1],
    #         [-2,0,2],
    #         [-1,0,1]]
    # ytmp = [[-1,-2,-1],
    #         [0,0,0],
    #         [1,2,1]]
    # xKernel = np.array(xtmp)
    # yKernel = np.array(ytmp)
    # gx = cv2.filter2D(image, -1, xKernel)
    # print(np.min(gx))
    # # cv2.imshow('gx',gx)
    # # cv2.waitKey(0)
    # gy = cv2.filter2D(image, -1, yKernel)
    # # cv2.imshow('gy',gy)
    # # cv2.waitKey(0)
    # magnitude = np.zeros_like(image, "float")
    # oritation = np.zeros_like(image, "float")
    edgeDetect = CannyEdgeDetector(image, kernel_type="sobel")
    gx, gy = edgeDetect.getGrads()
    magnitude, oritation = edgeDetect.getMagnitudeAndOrientation()
    # for i in range(h):
    #     for j in range(w):
    #         magnitude[i][j] = math.sqrt(gx[i][j] ** 2 + gy[i][j] ** 2) / 255
    #         # print(magnitude[i][j])
    #         oritation[i][j] = np.math.atan2(gy[i][j], gx[i][j])
    #         if oritation[i][j] < 0:
    #             oritation[i][j] += 2 * math.pi
            # print(oritation[i][j])
    # print(oritation)
    print("oritation:", np.max(oritation))
    print(oritation[83][19])
    # cv2.imshow('magnitude',magnitude / 255)
    # cv2.waitKey(0)
    # cv2.imshow('oritation',oritation)
    # cv2.waitKey(0)
    return magnitude, oritation

def getMainDirect(magnitude, oritation, keypoints, neighbours = 16, num_hi = 8):
    h = magnitude.shape[0]
    w = magnitude.shape[1]
    print(magnitude.shape)
    
    mainDirects = []
    for keypoint in keypoints:
        vote = [0 for i in range(num_hi)]
        for i in range(keypoint[1] - neighbours // 2, keypoint[1] + neighbours // 2):
            for j in range(keypoint[0] - neighbours // 2, keypoint[0] + neighbours // 2):
                vote[int(oritation[i][j] // (2 * math.pi / num_hi))] += 1
        # print(vote)
        maxx = 0
        p = -1
        for i in range(num_hi):
            if vote[i] > maxx:
                maxx = vote[i]
                p = i
        mainDirects.append(p * 2 * math.pi / num_hi)
    return mainDirects

def getDescriptors(magnitude, oritation, keypoints, mainDirects, neighbours = 16, num_hi = 8, cell_dim = 4):
    h = magnitude.shape[0]
    w = magnitude.shape[1]
    print(magnitude.shape)
    
    descriptors = []
    counter = 0
    for keypoint in keypoints:
        block = oritation[keypoint[1] - neighbours // 2 : keypoint[1] + neighbours // 2, keypoint[0] - neighbours // 2 : keypoint[0] + neighbours // 2]
        num_cell = neighbours // cell_dim
        startw = 0
        starth = 0
        descriptorList = []
        for i in range(num_cell):
            startw = 0
            for j in range(num_cell):
                vote = [0 for k in range(num_hi)]
                for m in range(starth, starth + cell_dim):
                    for n in range(startw, startw + cell_dim):
                        tmp = block[m][n] - mainDirects[counter]
                        if tmp < 0:
                            tmp += 2 * math.pi
                        vote[int(tmp // (2 * math.pi / num_hi))] += 1
                descriptorList.extend(vote)
                startw += cell_dim
            starth += cell_dim
        counter += 1
        descriptor = np.array(descriptorList)
        descriptors.append(descriptor)
    return descriptors


def keypointMatching(keypoints1, descriptors1, keypoints2, descriptors2):
    length1 = len(keypoints1)
    length2 = len(keypoints2)
    bestMatches = []
    for i in range(length1):
        bestMatch = [keypoints1[i]]
        best = -1
        bestSimilarity = 0
        for j in range(length2):
            similarity = getSimilarity(descriptors1[i], descriptors2[j])
            if similarity > bestSimilarity:
                bestSimilarity = similarity
                best = j
        bestMatch.append(keypoints2[best])
        bestMatches.append(bestMatch)
    return bestMatches

def getSimilarity(des1, des2):
    assert len(des1) == len(des2)
    sum_up = 0
    sum_1 = 0
    sum_2 = 0
    for i in range(len(des1)):
        sum_up += des1[i] * des2[i]
        sum_1 += des1[i] * des1[i]
        sum_2 += des2[i] * des2[i]
    return sum_up / sqrt(sum_1 * sum_2)

def getMatches():
    image1, keypoints1 = getKeyPoints("./SIFT/euclideanbackwardcat-0-0-0.5235987755982988.jpg")
    # print(np.max(image1))
    # print(keypoints[0])
    # print(len(keypoints))
    magnitude1, oritation1 = getGrad(image1)
    mainDirects1 = getMainDirect(magnitude1, oritation1, keypoints1, neighbours = 16, num_hi = 8)
    # print(mainDirects)
    descriptors1 = getDescriptors(magnitude1, oritation1, keypoints1, mainDirects1, neighbours = 16, num_hi = 8, cell_dim = 4)
    # print(len(descriptors1))
    # print(len(descriptors1[0]))
    # print(descriptors1[10])


    image2, keypoints2 = getKeyPoints("./SIFT/rawImage.jpg")
    # print(np.max(image1))
    # print(keypoints[0])
    # print(len(keypoints))
    magnitude2, oritation2 = getGrad(image2)
    mainDirects2 = getMainDirect(magnitude2, oritation2, keypoints2, neighbours = 16, num_hi = 8)
    # print(mainDirects)
    descriptors2 = getDescriptors(magnitude2, oritation2, keypoints2, mainDirects2, neighbours = 16, num_hi = 8, cell_dim = 4)
    # print(len(descriptors1))
    # print(len(descriptors1[0]))
    # print(descriptors1[10])
    matches = keypointMatching(keypoints1, descriptors1, keypoints2, descriptors2)
    print(len(matches))
    return matches

def judge(matches):
    forward_transformMatrix = np.array([[sqrt(3) / 2, - 0.5, 0], [0.5, sqrt(3) / 2, 0], [0, 0, 1]])
    transformMatrix = np.linalg.inv(forward_transformMatrix)
    avg = 0
    for item in matches:
        tmp = np.array([item[0][1] - 128, item[0][0] - 128, 1])
        fl = transformMatrix.dot(tmp)
        sum = abs(item[1][1] - fl[0] + item[1][0] - fl[1] - 2 * 128)
        avg += sum
    return avg / len(matches)


# if __name__ == "__main__":
#     image, keypoints = getKeyPoints("./SIFT/euclideanbackwardcat-0-0-0.5235987755982988.jpg")
#     print(np.max(image))
#     # print(keypoints[0])
#     # print(len(keypoints))
#     magnitude, oritation = getGrad(image)
#     mainDirects = getMainDirect(magnitude, oritation, keypoints, neighbours = 16, num_hi = 8)
#     # print(mainDirects)
#     descriptors = getDescriptors(magnitude, oritation, keypoints, mainDirects, neighbours = 16, num_hi = 8, cell_dim = 4)
#     print(len(descriptors))
#     print(len(descriptors[0]))
#     print(descriptors[10])

def getMatchingImg():
    rotateImg, _ = getKeyPoints("./SIFT/euclideanbackwardcat-0-0-0.5235987755982988.jpg")
    rotateImg = cv2.cvtColor(rotateImg, cv2.COLOR_GRAY2BGR)
    rawImage, __ = getKeyPoints("./SIFT/rawImage.jpg")
    rawImage = cv2.cvtColor(rawImage, cv2.COLOR_GRAY2BGR)
    matches = getMatches()
    h = 256
    w = 256 * 2
    img = np.zeros(shape = (h, w, 3), dtype = "float")
    img[:, :256] = rotateImg
    img[:, 256:] = rawImage
    counter = 0
    for item in matches:
        if counter > 10:
            break
        img = cv2.line(img, item[0], (item[1][0] + 256, item[1][1]), (0, 0, 255), 1)
        counter += 1
        
    cv2.imshow("Match", img)
    cv2.waitKey(0)
    cv2.imwrite("./SIFT/MatchingImage.jpg", img)


if __name__ == "__main__":
    # matches = getMatches()
    # print(judge(matches))
    # print(matches[0])
    # getMatchingImg()
    getKeyPoints("./SIFT/euclideanbackwardcat-0-0-0.5235987755982988.jpg")